from py2neo import Graph, Node
import pandas as pd
import os
from dotenv import load_dotenv
from rank_bm25 import BM25Okapi
import jieba

# ----------------------------
# 中文停用词（基础版，可替换为文件）
# ----------------------------
# 建议：后续可换成从文件加载，如 'stopwords.txt'
CHINESE_STOPWORDS = {
    '的', '了', '和', '是', '就', '都', '而', '及', '与', '或', '等', '在', '上', '下', '中', '内', '外',
    '一', '二', '三', '1', '2', '3', '个', '件', '种', '类', ' ', '\t', '\n', ''
}

# 可选：从文件加载停用词（取消注释并指定路径即可）
# with open("stopwords.txt", "r", encoding="utf-8") as f:
#     CHINESE_STOPWORDS = set(line.strip() for line in f if line.strip())

def preprocess_text(text):
    """使用 jieba 对中文文本分词，并去除停用词"""
    if not text or pd.isna(text):
        return []
    text = str(text).strip()
    if not text:
        return []
    # jieba 精确分词
    words = jieba.lcut(text)
    # 过滤停用词和空词
    filtered = [w.strip() for w in words if w.strip() and w not in CHINESE_STOPWORDS]
    return filtered

def node_to_dict_with_Id(node: Node):
    return {
        "identity": node.identity,
        "labels": list(node.labels),
        "properties": dict(node)
    }

# ----------------------------
# 1. 连接 Neo4j
# ----------------------------
load_dotenv()
NEO4J_URI = os.getenv("NEO4J_URI")
NEO4J_USER = os.getenv("NEO4J_USER")
NEO4J_PASSWORD = os.getenv("NEO4J_PASSWORD")
graph = Graph(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD))
print("✅ 连接Neo4j成功" if graph.run("RETURN 1").data() else "❌ 连接Neo4j失败")

# ----------------------------
# 2. 加载标准库
# ----------------------------
import pandas as pd
import os

# 工艺数据目录（请替换为实际路径）
data_dir = os.path.join("data_bm25_2")

df_list = []
for f in os.listdir(data_dir):
    if f.endswith('.csv'):
        # 读取文件并提取目标列
        df = pd.read_csv(os.path.join(data_dir, f))
        df = df[['编码', '名称', '规格型号']]
        # 添加"数据来源"列，记录当前文件名
        df['数据来源'] = f
        df_list.append(df)

# 合并成大DataFrame
df = pd.concat(df_list, ignore_index=True)
print(f"✅ 合并完成，共 {len(df)} 条记录")


# ----------------------------
# 3. 从图谱获取辅料节点
# ----------------------------
part_dict = [node_to_dict_with_Id(record["n"]) for record in graph.run("MATCH (n:`制造资源`) RETURN n")]
print(f"✅ 查询到 {len(part_dict)} 个辅料节点")

# ----------------------------
# 4. 构建 BM25 所需语料
# ----------------------------
df['combined_text'] = df.apply(
    lambda row: f"{row['名称']} {row['规格型号']}" 
    if pd.notna(row['名称']) and pd.notna(row['规格型号']) 
    else (row['名称'] if pd.notna(row['名称']) else str(row['规格型号'])),
    axis=1
)

corpus = [preprocess_text(text) for text in df['combined_text']]
bm25 = BM25Okapi(corpus)

# ----------------------------
# 5. 匹配图谱辅料到标准库
# ----------------------------
results = []
for part in part_dict:
    part_name = part['properties'].get('资源名称', '')
    if not part_name:
        continue
    query_tokens = preprocess_text(part_name)
    if not query_tokens:
        continue
    scores = bm25.get_scores(query_tokens)
    max_idx = scores.argmax()
    max_score = scores[max_idx]
    matched_row = df.iloc[max_idx]
    results.append({
        'neo4j_ID': part['identity'],
        'neo4j_名称': part_name,
        'standard_名称': matched_row['名称'],
        'standard_规格型号': matched_row['规格型号'],
        'standard_编码': matched_row.get('编码', ''),
        'standard_来源': matched_row.get('数据来源', ''),
        'max_score': max_score
    })

# ----------------------------
# 6. 输出结果
# ----------------------------
results_df = pd.DataFrame(results)
print(f"\n📊 共匹配 {len(results_df)} 个制造资源节点")
print(results_df.head(10))  # 只打印前10行，避免刷屏

# 定义来源类型映射
source_mapping = {
    '辅料': {'type': 'PartIteration', 'modelDefinition': 'Accessories'},
    '加工检测设备': {'type': 'ResourceIteration', 'modelDefinition': 'Equipment'},
    '加工工具': {'type': 'ResourceIteration', 'modelDefinition': 'Tool'},
    '工装工具': {'type': 'ResourceIteration', 'modelDefinition': 'Frock'},
    '计量器具': {'type': 'ResourceIteration', 'modelDefinition': 'MeasuringInstrument'}
}

# 添加两个新列
results_df['type'] = ''
results_df['modelDefinition'] = ''

# 遍历每一行，根据 standard_来源 填充新列
def fill_type_and_model(row):
    source = row['standard_来源']
    for key in source_mapping:
        if key in source:
            row['type'] = source_mapping[key]['type']
            row['modelDefinition'] = source_mapping[key]['modelDefinition']
            break
    return row

results_df = results_df.apply(fill_type_and_model, axis=1)


output_path = "out-制造资源匹配结果.csv"
results_df.to_csv(output_path, index=False, encoding='utf-8-sig')
print(f"\n✅ 结果已保存至：{output_path}")